package com.radar.algorithm;


import com.radar.common.GisConstants;
import org.apache.commons.math3.linear.MatrixUtils;
import org.apache.commons.math3.linear.RealMatrix;
import org.apache.commons.math3.stat.correlation.Covariance;
import org.bytedeco.javacpp.indexer.Indexer;
import org.bytedeco.opencv.global.opencv_imgcodecs;
import org.bytedeco.opencv.global.opencv_imgproc;
import org.bytedeco.opencv.opencv_core.*;
import org.meteoinfo.data.GridData;
import org.meteoinfo.geoprocess.analysis.ResampleMethods;
import org.meteoinfo.math.spatial.KDTree;
import org.apache.commons.math3.linear.LUDecomposition;
import wcontour.Interpolate;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.*;
import java.util.List;
import java.util.stream.IntStream;

import static org.bytedeco.opencv.global.opencv_core.*;
import static org.bytedeco.opencv.global.opencv_imgproc.*;
import static org.bytedeco.opencv.global.opencv_video.calcOpticalFlowPyrLK;

/**
 * @author: fy-wbj
 * @description: Interface to the OpenCV_ implementation of the local `Lucas-Kanade`_ optical
 * flow method applied in combination to a feature detection routine.
 * <p>
 * The sparse motion vectors are finally interpolated to return the whole
 * motion field.
 * @date: 2021/1/21
 */
public class DenseLucaskanade {

    public static double[][][] cacl(double[][][] data) {
        double[][][] input_images = clone(data);
        System.out.println("Computing the motion field with the Lucas-Kanade method.");
        long t0 = System.currentTimeMillis();
        int nr_fields = input_images.length;
        int[] domain_size = new int[]{input_images[0].length, input_images[0][0].length};
        List<double[]> xy = new ArrayList<>();
        List<double[]> uv = new ArrayList<>();

        for (int n = 0; n < data.length - 1; n++) {
            double[][] prvs_img = clone(input_images[n]);
            double[][] next_img = clone(input_images[n + 1]);
            // emove small noise with a morphological operator (opening)
            prvs_img = morph_opening(prvs_img);
            next_img = morph_opening(next_img);

            // features detection
            Mat points = shitomasiDetection(prvs_img);

            double[][][] doubles = track_features(prvs_img, next_img, points);
            double[][] xy_ = doubles[0];
            double[][] uv_ = doubles[1];
            // get sparse u, v vectors with Lucas-Kanade tracking
            for (int i = 0; i < xy_.length; i++) {
                xy.add(xy_[i]);
                uv.add(uv_[i]);
            }
        }
        // detect and remove outliers
        boolean[] booleans = detect_outliers(uv, 3, xy, 30);
        List<double[]> nxy = new ArrayList<>();
        List<double[]> nuv = new ArrayList<>();
        for (int i = 0; i < xy.size(); i++) {
            if (!booleans[i]) {
                nxy.add(xy.get(i));
                nuv.add(uv.get(i));
            }
        }
        xy = nxy;
        uv = nuv;

        List<List<double[]>> decluster = decluster(xy, uv);
        xy = decluster.get(0);
        uv = decluster.get(1);
        double[] xarray = new double[domain_size[1]];
        double[] yarray = new double[domain_size[0]];
        IntStream.range(0, xarray.length).forEach(i -> xarray[i] = i);
        IntStream.range(0, yarray.length).forEach(i -> yarray[i] = i);
        double[][] scoords = new double[xy.size()][3];
        double[][] scoords1 = new double[xy.size()][3];
        for (int i = 0; i < scoords.length; i++) {
            scoords[i][0] = xy.get(i)[0];
            scoords[i][1] = xy.get(i)[1];
            scoords[i][2] = uv.get(i)[0];

            scoords1[i][0] = xy.get(i)[0];
            scoords1[i][1] = xy.get(i)[1];
            scoords1[i][2] = uv.get(i)[1];
        }
        // 控制计算速度 插值到800 * 800 网格
        if (domain_size[1] > 800 || domain_size[0] > 800) {
            double xself = domain_size[1] / 800.0;
            double yself = domain_size[0] / 800.0;
            double[] nxarray = new double[(int) (domain_size[1] / xself) + 1];
            double[] nyarray = new double[(int) (domain_size[0] / yself) + 1];
            System.out.println("start interpolation_IDW_Neighbor u v " + nxarray.length + " * " + nyarray.length + " points");
            IntStream.range(0, nxarray.length).forEach(i -> nxarray[i] = i * xself);
            IntStream.range(0, nyarray.length).forEach(i -> nyarray[i] = i * yself);
            double[][] nu = Interpolate.interpolation_IDW_Neighbor(scoords, nxarray, nyarray, 4, 0);
            double[][] nv = Interpolate.interpolation_IDW_Neighbor(scoords1, nxarray, nyarray, 4, 0);
            GridData gridDatau = new GridData();
            gridDatau.xArray = nxarray;
            gridDatau.yArray = nyarray;
            gridDatau.projInfo = GisConstants.lonlat;
            gridDatau.data = nu;
            gridDatau.missingValue = 0;

            GridData gridDatav = (GridData) gridDatau.clone();
            gridDatav.data = nv;

            GridData toGridData = new GridData();
            toGridData.missingValue = 0;
            toGridData.projInfo = GisConstants.lonlat;
            toGridData.data = new double[yarray.length][xarray.length];
            toGridData.xArray = xarray;
            toGridData.yArray = yarray;
            GridData uresample = gridDatau.resample(toGridData, ResampleMethods.Bilinear);
            GridData vresample = gridDatav.resample(toGridData, ResampleMethods.Bilinear);
            double[][] u = uresample.data;
            double[][] v = vresample.data;
            System.out.println("end interpolation_IDW_Neighbor u v");
            return new double[][][]{u, v};
        } else {
            System.out.println("start interpolation_IDW_Neighbor u v");
            double[][] u = Interpolate.interpolation_IDW_Neighbor(scoords, xarray, yarray, 4, 0);
            double[][] v = Interpolate.interpolation_IDW_Neighbor(scoords1, xarray, yarray, 4, 0);
            System.out.println("end interpolation_IDW_Neighbor u v");
            return new double[][][]{u, v};
        }
    }

    /**
     * Detect outliers in a (multivariate and georeferenced) dataset.
     * <p>
     * Assume a (multivariate) Gaussian distribution and detect outliers based on
     * the number of standard deviations from the mean.
     * <p>
     * If spatial information is provided through coordinates, the outlier
     * detection can be localized by considering only the k-nearest neighbours
     * when computing the local mean and standard deviation.
     */
    public static boolean[] detect_outliers(List<double[]> input, int thr, List<double[]> xy, int k_) {
        double[][] input_array = input.toArray(new double[][]{});
        double[][] coord = xy.toArray(new double[][]{});
        int nsamples = coord.length;
        int nvar = 2;
        int k = k_ + 1;
        KDTree.Euclidean KD = new KDTree.Euclidean(2);
        for (int i = 0; i < nsamples; i++) {
            KD.addPoint(coord[i], i);
        }
        boolean[] outliers = new boolean[nsamples];
        for (int i = 0; i < nsamples; i++) {
            ArrayList<KDTree.SearchResult> arrayList = KD.nearestNeighbours(coord[i], k);
            double[] thisdata = new double[]{input_array[i][0], input_array[i][1]};
            double[][] neighbours = new double[k - 1][2];
            double z = 0;
            double z1 = 0;
            for (int j = 1; j < arrayList.size(); j++) {
                int index = (int) arrayList.get(j).payload;
                neighbours[j - 1] = new double[]{input_array[index][0], input_array[index][1]};
                z += neighbours[j - 1][0];
                z1 += neighbours[j - 1][1];
            }
            z = z / (k - 1);
            z1 = z1 / (k - 1);
            double[] thiszdata = new double[]{thisdata[0] - z, thisdata[1] - z1};
            for (int j = 0; j < neighbours.length; j++) {
                neighbours[j][0] = neighbours[j][0] - z;
                neighbours[j][1] = neighbours[j][1] - z1;
            }
            RealMatrix mx = MatrixUtils.createRealMatrix(neighbours);
            RealMatrix cov = new Covariance(mx).getCovarianceMatrix();
            double MD = 0;
            try {
                RealMatrix VI = new LUDecomposition(cov).getSolver().getInverse();
                RealMatrix mat1 = MatrixUtils.createRealMatrix(new double[][]{thiszdata});
                RealMatrix mat2 = mat1.multiply(VI);
                double c = mat1.getEntry(0, 0) * mat2.getEntry(0, 0) + mat2.getEntry(0, 1) * mat1.getEntry(0, 1);
                MD = Math.sqrt(c);
            } catch (Exception e) {
                MD = 0;
            }
            outliers[i] = MD > thr;
        }
        return outliers;
    }

    /**
     * Filter out small scale noise on the image by applying a binary
     * morphological opening, that is, erosion followed by dilation.
     *
     * @param in
     * @return
     */
    public static double[][] morph_opening(double[][] in) {
        double[][] intput_image = clone(in);
        // Build a structuring element of size n
        Mat kernel = getStructuringElement(2, new Size(3, 3));
        // Convert to binary image
        Mat field_bin = new Mat(new Size(intput_image[0].length, intput_image.length), CV_8UC1);
        Indexer data = field_bin.createIndexer();
        for (int i = 0; i < intput_image.length; i++) {
            for (int j = 0; j < intput_image[0].length; j++) {
                if (intput_image[i][j] > 0) {
                    data.putDouble(new long[]{i, j}, 1);
                } else {
                    data.putDouble(new long[]{i, j}, 0);
                }
            }
        }
        Mat field_bin_out = field_bin.clone();
        // Apply morphological opening (i.e. erosion then dilation)
        morphologyEx(field_bin, field_bin_out, MORPH_OPEN, kernel);
        Indexer indexer1 = field_bin.createIndexer();
        Indexer indexer = field_bin_out.createIndexer();

        for (int i = 0; i < intput_image.length; i++) {
            for (int j = 0; j < intput_image[0].length; j++) {
                if ((indexer1.getDouble(i, j) - indexer.getDouble(i, j)) > 0) {
                    intput_image[i][j] = 0;
                }
            }
        }
        return intput_image;
    }

    public static Mat shitomasiDetection(double[][] data) {
        double[][] intput_image = DenseLucaskanade.clone(data);
        Mat mat = new Mat(new Size(intput_image[0].length, intput_image.length), CV_8UC1);
        Mat mask = new Mat(new Size(intput_image[0].length, intput_image.length), CV_8UC1);
        Indexer indexer = mat.createIndexer();
        Indexer matIndexer = mask.createIndexer();
        double[] maxMinValue = new double[0];
        maxMinValue = MathCommon.getMaxMinValue(intput_image);
        for (int i = 0; i < intput_image.length; i++) {
            for (int j = 0; j < intput_image[0].length; j++) {
                // 替代 cv2.dilate 目前没有转换 TODO
                if (i == 0) {
                    intput_image[i][j] = 0;
                }
                if (intput_image[i][j] > 0) {
                    intput_image[i][j] = (int) (intput_image[i][j] / (maxMinValue[0]) * 255);
                }
                indexer.putDouble(new long[]{i, j}, Math.floor(intput_image[i][j]));
                matIndexer.putDouble(new long[]{i, j}, 1);
            }
        }
        Mat output = mat.clone();
        goodFeaturesToTrack(mat, output, 1000, 0.01, 10, mask, 5, false, 0.04);
        System.out.println("--- 1000 good features to track detected ---");
        return output;
    }

    /**
     * Interface to the OpenCV `Lucas-Kanade`_ features tracking algorithm
     * (cv.calcOpticalFlowPyrLK).
     *
     * @return
     */
    public static double[][][] track_features(double[][] pre, double[][] next, Mat points) {
        double[][] prvs_img = clone(pre);
        double[][] next_img = clone(next);
        Mat p = new Mat(new Size(prvs_img[0].length, prvs_img.length), CV_8U);
        Mat n = new Mat(new Size(prvs_img[0].length, prvs_img.length), CV_8U);
        Indexer indexerP = p.createIndexer();
        Indexer indexerN = n.createIndexer();

        double[] maxMinValue = new double[0];
        double[] maxMinValue1 = new double[0];
        maxMinValue = MathCommon.getMaxMinValue(prvs_img);
        maxMinValue1 = MathCommon.getMaxMinValue(next_img);

        for (int i = 0; i < prvs_img.length; i++) {
            for (int j = 0; j < prvs_img[0].length; j++) {
                if (prvs_img[i][j] > 0) {
                    prvs_img[i][j] = (int) (prvs_img[i][j] / (maxMinValue[0]) * 255);
                }
                if (next_img[i][j] > 0) {
                    next_img[i][j] = (int) (next_img[i][j] / (maxMinValue1[0]) * 255);
                }
                indexerP.putDouble(new long[]{i, j}, prvs_img[i][j]);
                indexerN.putDouble(new long[]{i, j}, next_img[i][j]);
            }
        }

        Mat p1 = new Mat(new Size(1000, 2), CV_32FC2);
        Mat st = new Mat(new Size(1000, 1), CV_32FC2);
        Mat _ = new Mat(new Size(1000, 2), CV_32FC2);
        Mat p0 = points.clone();
        calcOpticalFlowPyrLK(p, n, p0, p1, st, _, new Size(50, 50), 3, new TermCriteria(3, 10, 0), 0, 0.0001);
        // keep only features that have been found
        Indexer indexer0 = p0.createIndexer();
        Indexer indexer1 = p1.createIndexer();
        int size = p0.size().get(1);
        double[][] xy = new double[size][2];
        double[][] uv = new double[size][2];
        for (int i = 0; i < size; i++) {
            xy[i][0] = indexer0.getDouble(i, 0, 0);
            xy[i][1] = indexer0.getDouble(i, 0, 1);

            uv[i][0] = indexer1.getDouble(i, 0, 0) - xy[i][0];
            uv[i][1] = indexer1.getDouble(i, 0, 1) - xy[i][1];
        }
        System.out.println("--- 1000 sparse vectors found ---");
        return new double[][][]{xy, uv};
    }

    public static List<List<double[]>> decluster(List<double[]> xy, List<double[]> uv) {
        double[][] coord = xy.toArray(new double[][]{});
        double[][] input_array = uv.toArray(new double[][]{});
        double scale = 20.0;
        double[][] coord_ = new double[coord.length][coord[0].length];
        for (int i = 0; i < coord.length; i++) {
            for (int j = 0; j < coord[0].length; j++) {
                coord_[i][j] = Math.floor(coord[i][j] / scale);
            }
        }
        Map<String, List<Integer>> map = new HashMap<>();
        for (int i = 0; i < coord.length; i++) {
            String key = coord_[i][0] + "_" + coord_[i][1];
            List<Integer> list;
            if (map.containsKey(key)) {
                list = map.get(key);
                list.add(i);
            } else {
                list = new ArrayList<>();
                list.add(i);
                map.put(key, list);
            }
        }
        List<double[]> dcoord = new ArrayList<>();
        List<double[]> dinput = new ArrayList<>();
        map.forEach((k, v) -> {
            if (v.size() >= 1) {
                List<Double> list = new ArrayList<>();
                List<Double> list1 = new ArrayList<>();
                List<Double> list3 = new ArrayList<>();
                List<Double> list4 = new ArrayList<>();
                v.forEach(i -> {
                    list.add(coord[i][0]);
                    list1.add(coord[i][1]);
                    list3.add(input_array[i][0]);
                    list4.add(input_array[i][1]);
                });
                dcoord.add(new double[]{median(list), median(list1)});
                dinput.add(new double[]{median(list3), median(list4)});
            }
        });
        List<List<double[]>> list = new ArrayList();
        list.add(dcoord);
        list.add(dinput);
        return list;
    }

    public static double median(List<Double> total) {
        double j = 0;
        //集合排序
        Collections.sort(total);
        int size = total.size();
        if (size % 2 == 1) {
            j = total.get((size - 1) / 2);
        } else {
            j = (total.get(size / 2 - 1) + total.get(size / 2) + 0.0) / 2;
        }
        return j;
    }

    /**
     * Interface to the OpenCV `Shi-Tomasi`_ features detection method to detect
     * corners in an image.
     * <p>
     * Corners are used for local tracking methods.
     *
     * @param data
     * @return
     */
    public static double[][][] clone(double[][][] data) {
        double[][][] d = new double[data.length][data[0].length][data[0][0].length];
        for (int i = 0; i < d.length; i++) {
            for (int j = 0; j < d[0].length; j++) {
                for (int k = 0; k < d[0][0].length; k++) {
                    d[i][j][k] = data[i][j][k];
                }
            }
        }
        return d;
    }

    public static double[][] clone(double[][] data) {
        double[][] d = new double[data.length][data[0].length];
        for (int i = 0; i < d.length; i++) {
            for (int j = 0; j < d[0].length; j++) {
                d[i][j] = data[i][j];
            }
        }
        return d;
    }

    public static void main(String[] args) throws IOException {
        Mat ff1 = opencv_imgcodecs.imread("/Users/xx/Downloads/img_data/20201231/VOL/" + "20201231065000_VOL_LAYER_100_GRAY.png");
        Mat ff2 = opencv_imgcodecs.imread("/Users/xx/Downloads/img_data/20201231/VOL/" + "20201231070000_VOL_LAYER_100_GRAY.png");
        Mat prvs = ff1.clone();
        Mat next = ff1.clone();
        opencv_imgproc.cvtColor(ff1, prvs, opencv_imgproc.COLOR_BGR2GRAY);
        opencv_imgproc.cvtColor(ff2, next, opencv_imgproc.COLOR_BGR2GRAY);
        int w = prvs.size().get(0);
        int h = prvs.size().get(1);
        double[][] pre = new double[h][w];
        double[][] nex = new double[h][w];
        Indexer indexer = prvs.createIndexer();
        Indexer indexer1 = next.createIndexer();
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w; j++) {
                pre[i][j] = indexer.getDouble(i, j);
                if (pre[i][j] == 255) {
                    pre[i][j] = 0;
                }
                nex[i][j] = indexer1.getDouble(i, j);
                if (nex[i][j] == 255) {
                    nex[i][j] = 0;
                }
            }
        }
        double[][][] cacl = cacl(new double[][][]{pre, nex});
        double[][] u = cacl[0];
        double[][] v = cacl[1];
        SemiLagrangian semiLagrangian = new SemiLagrangian();
        float[][][] extrapolate = semiLagrangian.extrapolate(MathCommon.doubleToFloat(nex), MathCommon.doubleToFloat(u), MathCommon.doubleToFloat(v), 24);
        for (int i = 0; i < 24; i++) {
            BufferedImage bufferedImage = new BufferedImage(u[0].length, u.length, BufferedImage.TYPE_INT_ARGB);
            for (int j = 0; j < u.length; j++) {
                for (int k = 0; k < u[0].length; k++) {
                    int z = (int) (extrapolate[i][j][k] == Float.NaN ? 0 : extrapolate[i][j][k]);
                    bufferedImage.setRGB(k, j, new Color(z, z, z).getRGB());
                }
            }
            ImageIO.write(bufferedImage, "png", new File("/Users/xx/data/java/flow_" + i + ".png"));
        }
    }
}
